
#include "PsoModel.h"
#include "PsoParticle.h"
#include "RandomGsl.h"
#include "MathUtil.h"
#include "PsoInertiaUpdateMethods.h"
#include "PsoSharingMethods.h"
#include "PsoVelocityClampingMethods.h"

namespace jcrada {

    PsoModel::PsoModel( int size ) :
    _size( size ),
    _swarm( new PsoParticle*[size] ),
    _current_iteration( 0 ),
    _lower_fitness_better( true ),
    _random_search( new RandomGsl( ) ),
    _r1( new RandomGsl( ) ),
    _r2( new RandomGsl( ) ),
    _r3( new RandomGsl( ) ),
    _nostalgia( 2.0 ),
    _envy( 2.0 ),
    _sharing_method( NULL ),
    _velocity_clamping_method( NULL ),
    _inertia_update_method( NULL ) {
        for (int i = 0; i < getSize(); ++i) {
            setParticle(i, NULL);
        }
    }

    PsoModel::PsoModel( int size, const PsoParticle& particle, bool lower_fitness_better ) :
    _size( size ),
    _swarm( new PsoParticle*[size] ),
    _current_iteration( 0 ),
    _lower_fitness_better( lower_fitness_better ),
    _random_search( new RandomGsl( ) ),
    _r1( new RandomGsl( ) ),
    _r2( new RandomGsl( ) ),
    _r3( new RandomGsl( ) ),
    _nostalgia( 2.0 ),
    _envy( 2.0 ),
    _sharing_method( NULL ),
    _velocity_clamping_method( NULL ),
    _inertia_update_method( NULL ) {
        PsoParticle* clone = NULL;
        for (int i = 0; i < getSize(); ++i) {
            clone = particle.clone();
            initialize(*clone);
            setParticle(i, clone);
        }
    }

    PsoModel::~PsoModel( ) {
        for (int i = 0; i < getSize(); ++i) {
            delete _swarm[i];
        }
        delete [] _swarm;

        delete _random_search;
        delete _r1;
        delete _r2;
        delete _r3;
    }

    void PsoModel::initialize( PsoParticle& particle ) {
        for (int i = 0; i < particle.getSize(); ++i) {
            if (particle.getSpace(i) == PsoParticle::S_CONTINUOUS) {
                particle.setPosition(i, getRandomSearch().nextDouble(particle.getMinPosition(i),
                        particle.getMaxPosition(i)));
            } else if (particle.getSpace(i) == PsoParticle::S_DISCRETE) {
                particle.setPosition(i, getRandomSearch().nextBool());
            } else {
                assert(0);
            }
            particle.setBestNeighborhoodPosition(i, particle.getPosition(i));
            particle.setBestPersonalPosition(i, particle.getPosition(i));
        }
        double worst_fitness = isLowerFitnessBetter() ? 1.0 / 0.0 : -1.0 / 0.0;
        particle.setFitness(worst_fitness);
        particle.setFitnessForNeighborhoodBest(worst_fitness);
        particle.setFitnessForPersonalBest(worst_fitness);
    }

    void PsoModel::updateVelocity( PsoParticle& particle ) {
        double * position = particle.getPosition();
        double * pbest = particle.getBestPersonalPosition();
        double * nbest = particle.getBestNeighborhoodPosition();
        double * velocity = particle.getVelocity();
        double cognitive, social;
        for (int i = 0; i < particle.getSize(); ++i) {
            cognitive = getR1().nextDouble() * getNostalgia() * (pbest[i] - position[i]);
            social = getR2().nextDouble() * getEnvy() * (nbest[i] - position[i]);
            switch (particle.getSpace(i)) {
                case PsoParticle::S_CONTINUOUS :
                {
                    velocity[i] = (velocity[i] * particle.getInertia()) + cognitive + social;
                }
                    break;
                case PsoParticle::S_DISCRETE :
                {
                    velocity[i] = velocity[i] + cognitive + social;
                }
                    break;
                default:
                    BDEBUG(TO_STRING(particle.getSpace(i)));
                    assert(0);
            }
        }
    }

    void PsoModel::updatePosition( PsoParticle& particle ) {
        double update = 0.0;
        for (int i = 0; i < particle.getSize(); ++i) {
            update = particle.getPosition(i) + particle.getVelocity(i);
            switch (particle.getSpace(i)) {
                case PsoParticle::S_CONTINUOUS :
                {
                    particle.setPosition(i, update);
                    if (update > particle.getMaxPosition(i)) {
                        particle.setPosition(i, particle.getMaxPosition(i));
                    } else if (update < particle.getMinPosition(i)) {
                        particle.setPosition(i, particle.getMinPosition(i));
                    }
                }
                    break;
                case PsoParticle::S_DISCRETE :
                {
                    if (getR3().nextDouble() < Math::sigmoid(particle.getVelocity(i))) {
                        particle.setPosition(i, 1.0);
                    } else {
                        particle.setPosition(i, 0.0);
                    }
                }
                    break;
                default:
                    BDEBUG(TO_STRING(particle.getSpace(i)));
                    assert(0);
            }
        }
    }

    void PsoModel::updatePersonalPosition( PsoParticle& particle ) {
        if (isBetter(particle.getFitness(), particle.getFitnessForPersonalBest())) {
            for (int i = 0; i < particle.getSize(); ++i) {
                particle.setBestPersonalPosition(i, particle.getPosition(i));
            }
            particle.setFitnessForPersonalBest(particle.getFitness());
        }
    }

    void PsoModel::step( ) {
        for (int i = 0; i < getSize(); ++i) {
            updatePersonalPosition(getParticle(i));
        }
        getSharingMethod()->share(*this);
        PsoParticle* particle = NULL;
        for (int i = 0; i < getSize(); ++i) {
            particle = &getParticle(i);

            if (getInertiaUpdateMethod()) {
                getInertiaUpdateMethod()->updateInertia(*particle, *this);
            }

            updateVelocity(*particle);
            if (getVelocityClampingMethod()) {
                getVelocityClampingMethod()->clamp(*particle);
            }

            updatePosition(*particle);
        }
        setCurrentIteration(getCurrentIteration() + 1);
    }

    std::string PsoModel::toString( ) const {
        std::string result;
        for (int i = 0; i < getSize(); ++i) {
            result += getParticle(i).toString() + "\n";
        }
        return result;
    }

    void calculateFitnessF( PsoModel& model ) {
        double avg_fitness = 0.0;
        for (int i = 0; i < model.getSize(); ++i) {
            PsoParticle& particle = model.getParticle(i);
            double x, y, z;
            //            DEBUG("x: r=[" + TO_STRING(0.0) + "," + TO_STRING(gsl_pow_int(2, particle.getSize() / 2)) +
            //                    "]" + "; v=" + TO_STRING(MathUtil::binToInt(particle.getPosition(), particle.getSize() / 2)));
            x = Math::scale(0.0, gsl_pow_int(2, particle.getSize() / 2),
                    Math::binToInt(particle.getPosition(), particle.getSize() / 2),
                    -5.0, 5.0);
            //            DEBUG("y: r=[" + TO_STRING(0.0) + "," + TO_STRING(gsl_pow_int(2, particle.getSize() / 2)) +
            //                    "]" + "; v=" + TO_STRING(MathUtil::binToInt(particle.getPosition() + particle.getSize() / 2, particle.getSize() / 2)));
            y = Math::scale(0.0, gsl_pow_int(2, particle.getSize() / 2),
                    Math::binToInt(particle.getPosition() + particle.getSize() / 2, particle.getSize() / 2),
                    -5.0, 5.0);

            double r = pow(gsl_pow_2(x) + gsl_pow_2(y), 1.0 / 4.0);
            z = (sin(gsl_pow_2(x) + 3.0 * gsl_pow_2(y))) / (0.1 + gsl_pow_2(r));
            z += (gsl_pow_2(x) + 5.0 * gsl_pow_2(y)) * exp(1 - gsl_pow_2(r)) / 2.0;
            z += (gsl_pow_3(sin(2.0 * x)) + gsl_pow_3(cos(-2.0 * y))) / 0.5;
            particle.setFitness(z);
            avg_fitness += z;
        }
        BDEBUG("avg-fitness=" + TO_STRING(avg_fitness / model.getSize()));
    }

    void calculateFitness( PsoModel& model ) {
        double avg_fitness = 0.0;
        double inertia = 0.0;

        double avg_value = 0.0;
        for (int i = 0; i < model.getSize(); ++i) {
            PsoParticle& particle = model.getParticle(i);
            double fitness = 0.0;
            double avg_value_x = 0.0;
            for (int j = 0; j < particle.getSize(); ++j) {
                fitness += gsl_pow_2(particle.getPosition(j));
                avg_value_x += particle.getPosition(j);
            }
            avg_value += avg_value_x / particle.getSize();
            particle.setFitness(fitness);
            avg_fitness += fitness;
            inertia += particle.getInertia();
        }
        BDEBUG("avg-value=" + TO_STRING(avg_value / model.getSize()) + " avg-fitness=" + TO_STRING(avg_fitness / model.getSize()));
        //        DEBUG("avg-inertia=" + TO_STRING(inertia / model.getSize()));
    }
    ////////////////
    //XML
    ////////////////

    PsoModel* PsoModel::fromXml( const QDomElement& xml, bool* ok ) {
        int size = xml.attribute("size", "0").toInt(ok);
        PsoModel* result = new PsoModel(size);
        result->setCurrentIteration(xml.attribute("current-iteration", "0").toInt(ok));
        result->setLowerFitnessBetter(xml.attribute("lower-fitness-better", "1").toInt(ok));
        result->setNostalgia(xml.attribute("nostalgia", "2.0").toDouble(ok));
        result->setEnvy(xml.attribute("envy", "2.0").toDouble(ok));
        int method = xml.attribute("sharing-method", "-1").toInt(ok);
        switch (method) {
            case SM_NONE: result->setSharingMethod(NULL);
                break;
            case SM_RING: result->setSharingMethod(new RingSharingMethod());
                break;
            case SM_STAR: result->setSharingMethod(new StarSharingMethod());
                break;
            case SM_VON_NEUMANN: result->setSharingMethod(new VonNeumannSharingMethod());
                break;
            default:
                BDEBUG(TO_STRING(method));
                assert(0);
        }
        method = xml.attribute("inertia-update-method", "-1").toInt(ok);
        switch (method) {
            case IM_NONE: result->setInertiaUpdateMethod(NULL);
                break;
            case IM_LINEAR: result->setInertiaUpdateMethod(new LinearDecreasingInertia(100));
                BDEBUG("Inertia Methods NOT working correctly!!!")
                break;
            case IM_NON_LINEAR: result->setInertiaUpdateMethod(new NonLinearDecreasingInertia(100));
                BDEBUG("Inertia Methods NOT working correctly!!!")
                break;
            default:
                BDEBUG(TO_STRING(method));
                assert(0);
        }
        method = xml.attribute("velocity-clamping-method", "-1").toInt(ok);
        switch (method) {
            case VM_NONE: result->setVelocityClampingMethod(NULL);
                break;
            case VM_BASIC: result->setVelocityClampingMethod(new BasicVelocityClamping());
                break;
            case VM_TANH: result->setVelocityClampingMethod(new TanhVelocityClamping());
                break;
            default:
                BDEBUG(TO_STRING(method));
                assert(0);
        }
        QDomNode child = xml.firstChild();
        int i = 0;
        while (!child.isNull()) {
            QDomElement e = child.toElement();
            if (e.tagName() == "PsoParticle") {
                assert(i < size);
                result->setParticle(i++, PsoParticle::fromXml(e, ok));
            }
            child = child.nextSibling();
        }

        if (!*ok) {
            BDEBUG("Error in PsoModel::fromXml");
        }
        return result;
    }

    void PsoModel::toXml( QDomDocument& xml, QDomElement* parent ) const {
        QDomElement model = xml.createElement("PsoModel");
        model.setAttribute("size", getSize());
        model.setAttribute("current-iteration", getCurrentIteration());
        model.setAttribute("lower-fitness-better", isLowerFitnessBetter());
        model.setAttribute("nostalgia", getNostalgia());
        model.setAttribute("envy", getEnvy());
        int method = -1;
        if (getSharingMethod()) {
            if (dynamic_cast<RingSharingMethod*> (getSharingMethod())) {
                method = SM_RING;
            } else if (dynamic_cast<StarSharingMethod*> (getSharingMethod())) {
                method = SM_STAR;
            } else if (dynamic_cast<VonNeumannSharingMethod*> (getSharingMethod())) {
                method = SM_VON_NEUMANN;
            }
        }
        model.setAttribute("sharing-method", method);

        method = -1;
        if (getInertiaUpdateMethod()) {
            if (dynamic_cast<LinearDecreasingInertia*> (getInertiaUpdateMethod())) {
                method = IM_LINEAR;
            } else if (dynamic_cast<NonLinearDecreasingInertia*> (getInertiaUpdateMethod())) {
                method = IM_NON_LINEAR;
            }
        }
        model.setAttribute("inertia-update-method", method);

        method = -1;
        if (getVelocityClampingMethod()) {
            if (dynamic_cast<BasicVelocityClamping*> (getVelocityClampingMethod())) {
                method = VM_BASIC;
            } else if (dynamic_cast<TanhVelocityClamping*> (getVelocityClampingMethod())) {
                method = VM_TANH;
            }
        }
        model.setAttribute("velocity-clamping-method", method);
        for (int i = 0; i < getSize(); ++i) {
            getParticle(i).toXml(xml, &model);
        }

        if (parent) {
            parent->appendChild(model);
        } else {
            xml.appendChild(model);
        }
    }

    int PsoModel::main( int argc, char** argv ) {
        int particle_size = 100;
        int swarm_size = 100;
        int iterations = 5000;
        PsoParticle* particle = new PsoParticle(particle_size);
        particle->setMinPosition(0.0);
        particle->setMaxPosition(100.0);
        particle->setMaxVelocity(50.0);

        PsoModel model(swarm_size, *particle, true);
        model.setSharingMethod(new RingSharingMethod());
        model.setInertiaUpdateMethod(NULL);
        model.setVelocityClampingMethod(new BasicVelocityClamping());
        BDEBUG(model.toString());
        for (int i = 0; i < iterations; ++i) {
            calculateFitness(model);
            model.step();
        }
        return 0;
    }


}
